// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


/*
 * Copyright (c) 2024 by Rivos Inc.
 * Licensed under the Apache License, Version 2.0, see LICENSE for details.
 * SPDX-License-Identifier: Apache-2.0
 */

#pragma once

#include "breeze/utils/device_allocator.h"
#include "breeze/utils/device_vector.h"

namespace breeze {

template <typename T, typename Allocator = utils::device_allocator<T>>
class device_column : public utils::device_vector<T, Allocator> {
 public:
  explicit device_column(const Allocator& allocator = Allocator())
      : utils::device_vector<T, Allocator>(allocator) {}
  explicit device_column(utils::size_type size,
                         const Allocator& allocator = Allocator())
      : utils::device_vector<T, Allocator>(size, allocator) {}
};

template <typename T, typename Allocator = utils::device_allocator<T>>
class device_column_buffered {
  typedef typename Allocator::template rebind<T*>::other PtrAllocator;

 public:
  explicit device_column_buffered(const Allocator& allocator = Allocator())
      : buffers_{device_column<T, Allocator>(allocator),
                 device_column<T, Allocator>(allocator)},
        ptrs_(utils::device_vector<T*, PtrAllocator>(2,
                                                     PtrAllocator(allocator))) {
    UpdatePtrs();
  }
  explicit device_column_buffered(utils::size_type size,
                                  const Allocator& allocator = Allocator())
      : buffers_{device_column<T, Allocator>(size, allocator),
                 device_column<T, Allocator>(size, allocator)},
        ptrs_(utils::device_vector<T*, PtrAllocator>(2,
                                                     PtrAllocator(allocator))) {
    UpdatePtrs();
  }

  void resize(size_t size) {
    buffers_[0].resize(size);
    buffers_[1].resize(size);
    UpdatePtrs();
  }
  size_t size() const { return buffers_[0].size(); }
  device_column<T, Allocator>& buffer(int index) { return buffers_[index]; }
  utils::device_vector<T*, PtrAllocator>& ptrs() { return ptrs_; }

 private:
  void UpdatePtrs() {
    T* ptrs[] = {buffers_[0].data(), buffers_[1].data()};
    ptrs_.copy_from_host(ptrs, 2);
  }

  device_column<T, Allocator> buffers_[2];
  utils::device_vector<T*, PtrAllocator> ptrs_;
};

}  // namespace breeze
